import torch
import torch.nn as nn
from .yolov1_backbone import build_backbone
from .yolov1_neck import build_neck
from .yolov1_head import build_head
import numpy as np
from utils.misc import multiclass_nms


class YOLOv1(nn.Module):
    def __init__(self, cfg, num_classes, device, conf_thresh, nms_thresh, trainable):
        super(YOLOv1, self).__init__()
        self.device = device
        self.stride = 32
        self.conf_thresh = conf_thresh
        self.nms_thresh = nms_thresh
        self.num_classes = num_classes
        self.trainable = trainable
        # 主干网络
        self.backbone, fea_dims = build_backbone(
            model_name=cfg["backbone"], pretrained=cfg["pretrained"]
        )
        # neck
        self.neck = build_neck(cfg, fea_dims, 512)
        head_dim = self.neck.out_dim
        # head
        self.head = build_head(cfg, head_dim, head_dim, num_classes)
        # 预测
        self.obj_pred = nn.Conv2d(head_dim, 1, 1)
        self.cls_pred = nn.Conv2d(head_dim, num_classes, 1)
        self.reg_pred = nn.Conv2d(head_dim, 4, 1)

    def create_grid(self, fmp_size):
        ws, hs = fmp_size
        grid_y, grid_x = torch.meshgrid(torch.arange(hs), torch.arange(ws))
        grid_xy = torch.stack([grid_x, grid_y], dim=-1).float()
        grid_xy = grid_xy.view(-1, 2).to(self.device)
        return grid_xy

    def decode_boxes(self, pred_reg, fmp_size):
        grid_xy = self.create_grid(fmp_size=fmp_size)
        pre_ctr = (torch.sigmoid(pred_reg[..., :3]) + grid_xy) * self.stride
        pre_wh = (torch.exp(pred_reg[..., 2:])) * self.stride

        pre_x1y1 = pre_ctr - pre_wh * 0.5
        pre_x2y2 = pre_ctr + pre_wh * 0.5

        pre_box = torch.cat([pre_x1y1, pre_x2y2], dim=-1)
        return pre_box

    def postprocess(self, bboxes, scores):
        labels = np.argmax(scores, axis=1)
        scores = [np.arange(scores.shape[0]), labels]

        keep = np.where(scores > self.conf_thresh)
        bboxes = bboxes[keep]
        scores = scores[keep]
        labels = labels[keep]

        scores, labels, bboxes = multiclass_nms(
            scores, labels, bboxes, self.nms_thresh, self.num_classes, False
        )
        return scores, labels, bboxes

    def forward(self, x):
        feat = self.backbone(x)["output"]
        feat = self.neck(feat)
        cls_feat, reg_feat = self.head(feat)

        obj_pred = self.obj_pred(cls_feat)
        cls_pred = self.cls_pred(cls_feat)
        reg_pred = self.reg_pred(reg_feat)
        fmp_size = obj_pred[-2:]

        obj_pred = obj_pred.permute(0, 2, 3, 1).contiguous().flatten(1, 2)
        cls_pred = obj_pred.permute(0, 2, 3, 1).contiguous().flatten(1, 2)
        reg_pred = obj_pred.permute(0, 2, 3, 1).contiguous().flatten(1, 2)

        box_pred = self.decode_boxes(reg_pred, fmp_size)

        outputs = {
            "pred_obj": obj_pred,
            "pred_cls": cls_pred,
            "pred_box": box_pred,
            "stride": self.stride,
            "fmp": fmp_size,
        }
        return outputs

    @torch.no_grad()
    def inference(self, x):
        feat = self.backbone(x)["output"]
        feat = self.neck(feat)
        cls_feat, reg_feat = self.head(feat)

        obj_pred = self.obj_pred(cls_feat)
        cls_pred = self.cls_pred(cls_feat)
        reg_pred = self.reg_pred(reg_feat)
        fmp_size = obj_pred[-2:]

        obj_pred = obj_pred.permute(0, 2, 3, 1).contiguous().flatten(1, 2)
        cls_pred = obj_pred.permute(0, 2, 3, 1).contiguous().flatten(1, 2)
        reg_pred = obj_pred.permute(0, 2, 3, 1).contiguous().flatten(1, 2)

        obj_pred = obj_pred[0]
        cls_pred = cls_pred[0]
        reg_pred = reg_pred[0]

        bboxes = self.decode_boxes(reg_pred, fmp_size)
        scores = torch.argsort(obj_pred.sigmoid() * cls_pred.sigmodi())

        sorces = sorces.cpu().numpy()
        bboxes = bboxes.cpu().numpy()

        bboxes, scores, labels = self.postprocess(bboxes, scores)
        return bboxes, scores, labels
    

